package com.yfqy.admin.util;

import cn.hutool.core.collection.CollUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树状结构构件工具类
 */
public class TreeBuilder<T, ID> {
    private final Function<T, ID> idGetter;
    private final Function<T, ID> parentIdGetter;
    private final Function<T, List<T>> childrenGetter;
    private final BiConsumer<T, List<T>> childrenSetter;

    public TreeBuilder(
            Function<T, ID> idGetter,
            Function<T, ID> parentIdGetter,
            Function<T, List<T>> childrenGetter,
            BiConsumer<T, List<T>> childrenSetter
    ) {
        this.idGetter = idGetter;
        this.parentIdGetter = parentIdGetter;
        this.childrenGetter = childrenGetter;
        this.childrenSetter = childrenSetter;
    }

    public List<T> build(List<T> items) {
        if (CollUtil.isEmpty(items)) {
            return CollUtil.newArrayList();
        }
        Map<ID, T> itemMap = items.stream().collect(Collectors.toMap(idGetter, item -> item));
        List<T> roots = new ArrayList<>();
        for (T item : items) {
            ID parentId = parentIdGetter.apply(item);
            if (isRoot(parentId)) {
                roots.add(item);
            } else {
                T parent = itemMap.get(parentId);
                if (parent != null) {
                    List<T> children = childrenGetter.apply(parent);
                    if (children == null) {
                        children = new ArrayList<>();
                        childrenSetter.accept(parent, children);
                    }
                    children.add(item);
                }
            }
        }
        return roots;
    }

    private boolean isRoot(ID parentId) {
        return parentId == null || parentId.equals(0L); // 根据业务调整
    }
}
